# -*- coding: utf-8 -*-
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks


class ModelPredict(object):

    def __init__(self):
        self.model = pipeline(Tasks.text_generation, model='damo/nlp_palm2.0_text-generation_chinese-base', model_revision='v1.0.1')

    def predict(self, text):
        preds = self.model(text)
        # return preds['text']
        return {
            "text": preds['text']
        }
